# python3
from bs4 import BeautifulSoup
import requests
import json
import arrow
import tushare as ts
from tqdm import tqdm
import time
import tempfile
import csv
import sys
sys.setrecursionlimit(10000) 
import logging
from functools import wraps
from collections import defaultdict
from config import get_header

def get_log():
    log = logging.getLogger(__name__)
    ft = logging.Formatter('%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s','%Y-%m-%d %H:%M:%S')

    fh = logging.FileHandler('log/%s.log'%__name__)
    fh.setFormatter(ft)
    fh.setLevel(logging.WARNING)

    sh = logging.StreamHandler()
    sh.setFormatter(ft)
    sh.setLevel(logging.DEBUG)

    log.addHandler(fh)
    log.addHandler(sh)

    log.setLevel(logging.DEBUG)
    return log

log = get_log()

class setLog:
    @classmethod
    def request(cls, retry=3):
        def decorate(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
                exception_count = defaultdict(lambda:0)
                for _ in range(retry):
                    try:
                        return func(*args, **kwargs)
                    except requests.exceptions.RequestException as e:
                        exception_count['RequestException'] += 1
                        log.warning(e)
                        time.sleep(5)
                    except AssertionError as e:
                        exception_count['AssertionError'] += 1
                        log.warning('%s %s'%(str(e),str(args if args else kwargs)))
                        if exception_count['AssertionError'] >= 1:
                            log.error('%s %s time over'%(str(e),str(args if args else kwargs)))
                            return
                        time.sleep(5)
                    except TypeError as e:
                        exception_count['TypeError'] += 1
                        log.warning('%s %s'%(str(e),str(args if args else kwargs)))
                        if exception_count['TypeError'] >= 1:
                            log.error('%s %s TypeError'%(str(e),str(args if args else kwargs)))
                            return
                        time.sleep(1)
                    except Exception as e:
                        exception_count['Exception'] += 1
                        log.exception(e)
                        time.sleep(5)
                else:
                    log.error('request fail %s'%str((args if args else kwargs)))
            return wrapper
        return decorate


class stock:
    def parse_dic_data(self, dic, symbol, arrnow, arrupdate=None):
        dic['_id'] = symbol + '_' + arrnow.format('YYYYMMDD')
        dic['symbol'] = symbol
        dic['update'] = arrupdate.datetime if arrupdate else arrnow.datetime

    def __repr__(self):
        return self.__class__.__name__


class Snowball(stock):

    def __init__(self):
        self.session = requests.session()
        self.session.headers.update(get_header())
        self.session.get('https://xueqiu.com',timeout=10)
        self.us_stocks = []

    def get_all_us_stock(self):
        def parse_org(dic):
            self.parse_dic_data(dic,dic.get('symbol'),arrow.now().floor('hour'))
            return dic
        self._get_all_us_stock()

        self.us_stocks = list(map(parse_org, self.us_stocks))
        return self.us_stocks

    @setLog.request()
    def _get_all_us_stock(self,page=1, total=None):
        '''
            https://xueqiu.com/hq#exchange=US&firstName=3&secondName=3_0
        '''
        url = 'https://xueqiu.com/stock/cata/stocklist.json?page={page}&size=90&order=desc&orderby=percent&type=0%2C1%2C2%2C3&isdelay=1'
        
        req = self.session.get(url.format_map({'page':page}),timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'
        req_json = req.json()
        stocks = req_json.get('stocks')
        self.us_stocks.extend(stocks)

        print(len(self.us_stocks), end='\r')
        if not total:
            total = req_json['count']['count']
        if page*90 >= total:
            return
        else:
            page += 1
            self._get_all_us_stock(page,total)


    def get_all_hk_stock(self):
        def parse_org(dic):
            self.parse_dic_data(dic,dic.get('symbol'),arrow.now().floor('hour'))
            return dic
        self._get_all_hk_stock()

        self.us_stocks = list(map(parse_org, self.us_stocks))
        return self.us_stocks

    @setLog.request()
    def get_hk_des(self,symbol):
        '''
            https://xueqiu.com/S/00147
        '''        
        def format_html(html):
            soup = BeautifulSoup(html,'lxml')
            desc = soup.select_one('.stock-company .detailContent')
            if not desc:
                desc = soup.select_one('.stock-company .summaryContent')
            info = soup.find(text='业务:')
            ret_dic = {}
            if desc:
                ret_dic['desc'] = desc.get_text().lstrip('简介: ')
            if info:
                ret_dic['info'] = info.next_element.strip()
            return ret_dic

        url = 'https://xueqiu.com/S/{symbol}'
        
        req = self.session.get(url.format_map({'symbol':symbol}),timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'

        ret_dic = format_html(req.text)

        return ret_dic

    @setLog.request()
    def _get_all_hk_stock(self,page=1, total=None):
        '''
            https://xueqiu.com/hq#exchange=US&firstName=3&secondName=3_0
        '''
        url = 'https://xueqiu.com/stock/cata/stocklist.json?page={page}&size=90&order=desc&orderby=percent&type=30&isdelay=1'
        
        req = self.session.get(url.format_map({'page':page}),timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'
        req_json = req.json()
        stocks = req_json.get('stocks')
        self.us_stocks.extend(stocks)

        print(len(self.us_stocks), end='\r')
        if not total:
            total = req_json['count']['count']
        if page*90 >= total:
            return
        else:
            page += 1
            self._get_all_hk_stock(page,total)


    @setLog.request()
    def get_us_finance(self, symbol):
        '''
        https://xueqiu.com/S/DMTX/keyratios
        '''
        def parse_date(dic):
            if 'TTM' == dic.get('date'):
                exrightdate = arrow.get('2100-01-01').replace(tzinfo='local')
            else:
                exrightdate = arrow.get(dic.get('date'),'YYYY-MM').replace(tzinfo='local')
            self.parse_dic_data(dic,symbol,exrightdate,arrow.now().floor('hour'))
            dic['exrightdate'] = exrightdate.datetime
            return dic

        url = 'https://xueqiu.com/stock/finance_us_financials.json'
        req = self.session.get(url,params={'symbol':symbol.replace('+','%2B')},timeout=10) 
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'
        lis = req.json().get('list')

        lis = list(map(parse_date,lis)) 

        # assert lis ,'ret lis is None'
        return lis

    @setLog.request()
    def get_oversea_quote(self,symbol):
        '''
        https://xueqiu.com/S/DMTX
        '''
        url = 'https://xueqiu.com/v4/stock/quote.json?code={symbol}&isUSdelay=1'
        req = self.session.get(url.format_map({'symbol':symbol.replace('+','%2B')}),timeout=10) 
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'
        dic = req.json().get(symbol)
        self.parse_dic_data(dic,symbol,arrow.now().floor('hour'))
        return dic


    @setLog.request()
    def get_h_stock(self,symbol,typ):
        '''
        https://xueqiu.com/S/OC430002
        typ: oc is neeq
        '''
        url = 'https://xueqiu.com/stock/forchartk/stocklist.json'
        req = self.session.get(url,params={'symbol':'%s%s'%(typ,symbol.replace('+','%2B')),'period':'1day','type':'normal'},timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'
        lis = req.json().get('chartlist')
        def parse_dic(dic):
            begindate = arrow.get(dic.get('time'),'MMM D HH:mm:ss Z YYYY')
            self.parse_dic_data(dic,symbol,begindate, arrow.now().floor('hour'))
            dic['time'] = begindate.datetime
            return dic

        lis = list(map(parse_dic,lis))
        assert lis ,'ret lis is None'
        return lis


import pymongo
Client = pymongo.MongoClient('192.168.1.220:29001')

def get_us():
    DB = Client.us_stock

    version = 1

    sbl = Snowball()

    if DB['index'].find_one({'version':version}):
        symbols = [symbol.get('symbol') for symbol in DB['index'].find({'version':version},{'symbol':1})]
    else:

        for com in sbl.get_all_us_stock():

            try:
                his_id = com['_id']
                com['_id'] = com['symbol']
                com['version'] = version

                DB['index'].delete_one({'_id':com['_id']})
                DB['index'].insert_one(com)
                
                com['_id'] = his_id
                DB['index_history'].insert_one(com)
            except pymongo.errors.DuplicateKeyError:
                pass

    # less
    # finish = [symbol.get('symbol') for symbol in DB['finance_us_financials'].find({'version':version},{'symbol':1})]
    # fi_symbols = set(symbols) - set(finish)
    # for symbol in tqdm(fi_symbols,desc='finance_us_financials'):
    #     data = sbl.get_us_finance(symbol)
    #     if data:
    #         for com in data:
    #             try:
    #                 com['version'] = version
    #                 DB['finance_us_financials'].update_one({'_id':com['_id']},{'$set':com},upsert=True)
    #             except pymongo.errors.DuplicateKeyError:
    #                 # repeat insert
    #                 pass


    finish = [symbol.get('symbol') for symbol in DB['quote'].find({'version':version},{'symbol':1})]
    fi_symbols = set(symbols) - set(finish)
    for symbol in tqdm(fi_symbols,desc='quote'):
        com = sbl.get_oversea_quote(symbol)
        if com:

            try:
                his_id = com['_id']
                com['_id'] = com['symbol']
                com['version'] = version

                DB['quote'].delete_one({'_id':com['_id']})
                DB['quote'].insert_one(com)
                
                com['_id'] = his_id
                DB['quote_history'].insert_one(com)
            except pymongo.errors.DuplicateKeyError:
                pass
    #     time.sleep(0.5)
    #     # break


    finish = [symbol.get('symbol') for symbol in DB['history_stock'].find({'version':version},{'symbol':1})]
    hi_symbols = set(symbols) - set(finish)
    for symbol in tqdm(hi_symbols,desc='get_h_stock'):
        data = sbl.get_h_stock(symbol, '')
        try:
            if data:
                list(map(lambda x:x.update({'version':version}),data))
                DB['history_stock'].insert_many(data,ordered=False)
        except pymongo.errors.BulkWriteError:
            # repeat insert
            pass
        time.sleep(0.5)

def get_hk():
    DB = Client.hk_stock

    version = 1

    sbl = Snowball()

    if DB['index'].find_one({'version':version}):
        symbols = [symbol.get('symbol') for symbol in DB['index'].find({'version':version},{'symbol':1})]
    else:

        for com in sbl.get_all_hk_stock():

            try:
                his_id = com['_id']
                com['_id'] = com['symbol']
                com['version'] = version

                DB['index'].delete_one({'_id':com['_id']})
                DB['index'].insert_one(com)
                
                com['_id'] = his_id
                DB['index_history'].insert_one(com)
            except pymongo.errors.DuplicateKeyError:
                pass



    finish = [symbol.get('_id') for symbol in DB['description'].find({'version':version},{'_id':1})]
    fi_symbols = set(symbols) - set(finish)
    for symbol in tqdm(fi_symbols,desc='description'):
        com = sbl.get_hk_des(symbol)
        if com:

            try:
                com['_id'] = symbol
                com['version'] = version
                DB['description'].insert_one(com)
                
            except pymongo.errors.DuplicateKeyError:
                pass
        # break


    finish = [symbol.get('symbol') for symbol in DB['history_stock'].find({'version':version},{'symbol':1})]
    hi_symbols = set(symbols) - set(finish)
    for symbol in tqdm(hi_symbols,desc='get_h_stock'):
        data = sbl.get_h_stock(symbol, '')
        try:
            if data:
                list(map(lambda x:x.update({'version':version}),data))
                DB['history_stock'].insert_many(data,ordered=False)
        except pymongo.errors.BulkWriteError:
            # repeat insert
            pass
        # time.sleep(0.)
        # break
if __name__ == '__main__':
    # get_hk()
    get_us()